import os
import threading
import time
from flask import Flask, render_template, request
from flask_socketio import SocketIO, emit, disconnect
import paramiko
import select
import socket
from threading import Thread
import logging

class TerminalBackend:
    def __init__(self, app_name='terminal_app', host='0.0.0.0', port=5000, debug=True):
        """
        初始化终端后端服务
        :param app_name: Flask应用名称
        :param host: 服务器地址
        :param port: 服务器端口
        :param debug: 调试模式
        """
        self.app = Flask(app_name)
        self.app.config['SECRET_KEY'] = 'your-secret-key-here'
        self.socketio = SocketIO(self.app, cors_allowed_origins="*")
        
        # 存储每个客户端的SSH连接
        self.ssh_connections = {}
        self.ssh_shells = {}
        
        # 服务器配置
        self.host = host
        self.port = port
        self.debug = debug
        
        # 注册事件处理器
        self._register_handlers()
        
        # 配置日志
        logging.basicConfig(level=logging.INFO)
        self.logger = logging.getLogger(__name__)

    def _register_handlers(self):
        """注册SocketIO事件处理器"""
        
        @self.socketio.on('connect')
        def handle_connect():
            """客户端连接事件"""
            client_id = request.sid
            self.logger.info(f"Client {client_id} connected")
            emit('terminal_output', {'data': '=== 终端已连接 ===\n请先使用 ssh_connect 命令连接到服务器\n'})

        @self.socketio.on('disconnect')
        def handle_disconnect():
            """客户端断开连接事件"""
            client_id = request.sid
            self.logger.info(f"Client {client_id} disconnected")
            self._cleanup_ssh_connection(client_id)

        @self.socketio.on('terminal_command')
        def handle_command(data):
            """处理终端命令"""
            client_id = request.sid
            command = data.get('cmd', '').strip()
            
            if not command:
                emit('terminal_output', {'data': ''})
                return
            
            try:
                # 检查是否是特殊命令
                if command.startswith('ssh_connect'):
                    self._handle_ssh_connect(client_id, command)
                elif command == 'ssh_disconnect':
                    self._handle_ssh_disconnect(client_id)
                elif command == 'clear':
                    emit('terminal_clear', {})
                elif command == 'help':
                    self._send_help(client_id)
                else:
                    # 执行SSH命令
                    self._execute_ssh_command(client_id, command)
                    
            except Exception as e:
                self.logger.error(f"Command execution error: {e}")
                emit('terminal_output', {'data': f'错误: {str(e)}\n'})

    def _handle_ssh_connect(self, client_id, command):
        """处理SSH连接命令"""
        try:
            # 解析连接参数: ssh_connect user@host:port password
            parts = command.split()
            if len(parts) < 3:
                emit('terminal_output', {
                    'data': '用法: ssh_connect user@host:port password\n例如: ssh_connect root@192.168.1.100:22 mypassword\n'
                })
                return
            
            user_host_port = parts[1]
            password = parts[2]
            
            # 解析用户名、主机和端口
            if '@' not in user_host_port:
                emit('terminal_output', {'data': '格式错误: 请使用 user@host:port 格式\n'})
                return
            
            user, host_port = user_host_port.split('@', 1)
            if ':' in host_port:
                host, port = host_port.split(':', 1)
                port = int(port)
            else:
                host = host_port
                port = 22
            
            # 建立SSH连接
            emit('terminal_output', {'data': f'正在连接到 {user}@{host}:{port}...\n'})
            
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh.connect(hostname=host, port=port, username=user, password=password, timeout=10)
            
            # 创建交互式shell
            shell = ssh.invoke_shell()
            shell.settimeout(0.1)
            
            # 存储连接
            self.ssh_connections[client_id] = ssh
            self.ssh_shells[client_id] = shell
            
            emit('terminal_output', {'data': f'成功连接到 {user}@{host}:{port}\n'})
            
            # 启动输出监听线程
            self._start_shell_listener(client_id, shell)
            
        except Exception as e:
            self.logger.error(f"SSH connection error: {e}")
            emit('terminal_output', {'data': f'连接失败: {str(e)}\n'})

    def _handle_ssh_disconnect(self, client_id):
        """处理SSH断开连接"""
        if client_id in self.ssh_connections:
            self._cleanup_ssh_connection(client_id)
            emit('terminal_output', {'data': 'SSH连接已断开\n'})
        else:
            emit('terminal_output', {'data': '没有活动的SSH连接\n'})

    def _execute_ssh_command(self, client_id, command):
        """执行SSH命令"""
        if client_id not in self.ssh_shells:
            emit('terminal_output', {'data': '请先建立SSH连接 (使用 ssh_connect 命令)\n'})
            return
        
        try:
            shell = self.ssh_shells[client_id]
            # 发送命令
            shell.send(command + '\n')
            
        except Exception as e:
            self.logger.error(f"Command execution error: {e}")
            emit('terminal_output', {'data': f'命令执行错误: {str(e)}\n'})

    def _start_shell_listener(self, client_id, shell):
        """启动shell输出监听线程"""
        def listen():
            while client_id in self.ssh_shells:
                try:
                    if shell.recv_ready():
                        output = shell.recv(4096).decode('utf-8', errors='ignore')
                        self.socketio.emit('terminal_output', {'data': output}, room=client_id)
                    time.sleep(0.01)
                except Exception as e:
                    if client_id in self.ssh_shells:
                        self.logger.error(f"Shell listener error: {e}")
                        self.socketio.emit('terminal_output', {'data': f'连接中断: {str(e)}\n'}, room=client_id)
                        self._cleanup_ssh_connection(client_id)
                    break
        
        thread = Thread(target=listen, daemon=True)
        thread.start()

    def _cleanup_ssh_connection(self, client_id):
        """清理SSH连接"""
        if client_id in self.ssh_shells:
            try:
                self.ssh_shells[client_id].close()
            except:
                pass
            del self.ssh_shells[client_id]
        
        if client_id in self.ssh_connections:
            try:
                self.ssh_connections[client_id].close()
            except:
                pass
            del self.ssh_connections[client_id]

    def _send_help(self, client_id):
        """发送帮助信息"""
        help_text = """
=== 终端帮助 ===
ssh_connect user@host:port password  - 连接到SSH服务器
ssh_disconnect                       - 断开SSH连接
clear                               - 清空终端
help                                - 显示此帮助信息

连接后可以执行任何Linux命令，如：
ls, cd, pwd, cat, vi, top 等

例子:
ssh_connect root@192.168.1.100:22 mypassword
"""
        emit('terminal_output', {'data': help_text})

    def add_route(self, rule, endpoint=None, view_func=None, **options):
        """添加自定义路由"""
        return self.app.route(rule, endpoint=endpoint, **options)(view_func)

    def run(self):
        """启动服务器"""
        self.logger.info(f"Starting terminal server on {self.host}:{self.port}")
        self.socketio.run(self.app, host=self.host, port=self.port, debug=self.debug)

# 使用示例
if __name__ == '__main__':
    # 创建终端后端实例
    terminal = TerminalBackend(host='0.0.0.0', port=5000)
    
    # 添加主页路由（可选）
    @terminal.add_route('/')
    def index():
        return """
        <!DOCTYPE html>
        <html>
        <head>
            <title>Terminal Backend</title>
        </head>
        <body>
            <h1>Terminal Backend is running!</h1>
            <p>Use your frontend to connect via SocketIO</p>
        </body>
        </html>
        """
    
    # 启动服务器
    terminal.run()